"""随机性参数,红酒数据"""
from sklearn import tree
from sklearn.datasets import load_wine
from sklearn.model_selection import train_test_split
import pandas as pd
import graphviz

if __name__ == "__main__":
    wine = load_wine()
    # print(wine.get("data"))  # 红酒的数据
    # print(wine.get("target"))
    print(wine.feature_names)
    table = pd.concat([pd.DataFrame(wine.get("data")), pd.DataFrame(wine.get("target"))], axis=1)  # 将数据转换成表
    # print(table)
    # 将数据分成训练集和测试集, test_size表示 30% 做测试集
    x_train, x_test, y_train, y_test = train_test_split(wine.data, wine.target, test_size=0.3)
    clf = tree.DecisionTreeClassifier(criterion="entropy")
    # 进行训练
    clf.fit(x_train, y_train)
    # 返回预测的准确度
    score = clf.score(x_test, y_test)
    print(score)

    # 将表画成树
    feature_name = ['酒精', '苹果酸', '灰', '灰的碱性', '镁', '总酚', '类黄酮', '非黄烷类酚类', '花青素', '颜色强度', '色调', 'od280 / od315稀释葡萄酒',
                    '脯氨酸']
    dot_data = tree.export_graphviz(
        clf,
        feature_names=feature_name,
        class_names=["琴酒", "雪莉", "贝尔摩德"],
        filled=True,  # 是否填充颜色
        rounded=True  # 方框
    )
    # 画图
    graph = graphviz.Source(dot_data)
    print(graph)
    # 查看重要的特征
    go = [*zip(feature_name, clf.feature_importances_)]
    print(go)

